package com.boge.tools;

import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import dev.langchain4j.service.tool.ToolExecutor;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

public class MainDemo1 {
    public static void main(String[] args) {
        // 获取 OpenAiChatModel 对象
        String apiKey = System.getenv("AI-BAILIAN-API-KEY");
        OpenAiChatModel model = OpenAiChatModel.builder()
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .apiKey(apiKey)
                .modelName("deepseek-r1")
                .logRequests(true)
                .logResponses(true)
                .build();

        // STEP 1: User specify tools and query
        // Tools
        WeatherTools weatherTools = new WeatherTools();
        List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(weatherTools);
        // User query
        List<ChatMessage> chatMessages = new ArrayList<>();
        UserMessage userMessage = new UserMessage("北京明天的天气怎么样?");
        chatMessages.add(userMessage);
        // Chat request
        ChatRequest chatRequest = ChatRequest.builder()
                .messages(chatMessages)
                .parameters(ChatRequestParameters.builder()
                        .toolSpecifications(toolSpecifications)
                        .build())
                .build();


        // STEP 2: Model generates tool execution request
        ChatResponse chatResponse = model.chat(chatRequest);
        AiMessage aiMessage = chatResponse.aiMessage();
        List<ToolExecutionRequest> toolExecutionRequests = aiMessage.toolExecutionRequests();
        System.out.println("Out of the " + toolSpecifications.size() + " tools declared in WeatherTools, " + toolExecutionRequests.size() + " will be invoked:");
        toolExecutionRequests.forEach(toolExecutionRequest -> {
            System.out.println("Tool name: " + toolExecutionRequest.name());
            System.out.println("Tool args:" + toolExecutionRequest.arguments());
        });
        chatMessages.add(aiMessage);


        // STEP 3: User executes tool(s) to obtain tool results
        toolExecutionRequests.forEach(toolExecutionRequest -> {
            ToolExecutor toolExecutor = new DefaultToolExecutor(weatherTools, toolExecutionRequest);
            System.out.println("Now let's execute the tool " + toolExecutionRequest.name());
            String result = toolExecutor.execute(toolExecutionRequest, UUID.randomUUID().toString());
            ToolExecutionResultMessage toolExecutionResultMessages = ToolExecutionResultMessage.from(toolExecutionRequest, result);
            chatMessages.add(toolExecutionResultMessages);
        });


        // STEP 4: Model generates final response
        ChatRequest chatRequest2 = ChatRequest.builder()
                .messages(chatMessages)
                .parameters(ChatRequestParameters.builder()
                        .toolSpecifications(toolSpecifications)
                        .build())
                .build();
        ChatResponse finalChatResponse = model.chat(chatRequest2);
        System.out.println(finalChatResponse.aiMessage().text()); //According to the payment data, the payment status of transaction T1005 is Pending.
    }
}
